{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  OboeLang.TopSyntax
-- Copyright   :  (c) Stephen Tetley 2014
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Top level syntax.
--
--------------------------------------------------------------------------------

module OboeLang.TopSyntax
  (

  -- * Syntax
    Program(..)
  , ModuleDict
  , Module(..)
  , ImportDecl(..)
  , OrchDict
  , FormExpr(..)
  , Binding(..)
  , BindingRhs(..)
  , DoBlock
  , Stmt(..)
  , VarBind(..)
  , Decl(..)
  , Expr(..)
  , Literal(..)

  -- * Pretty print  
  , ppProgram
  , ppModule

  -- * Operations
  , moduleDict
  , findModule

  ) where

import OboeLang.SyntaxBase
import OboeLang.Utils


import Text.PrettyExprHPJ                       -- package: pretty-expr-hpj
import Text.PrettyExprHPJ.Haskell98

import Text.PrettyPrint.HughesPJ                -- package: pretty

import qualified Data.IntMap as IntMap
import qualified Data.Map as Map

data Program = Program
     { prog_main                :: Module
     , prog_extra_libs          :: [Module]
     }
  deriving (Eq,Show)

type ModuleDict = Map.Map ModuleName Module

-- Only main should have an OrchDict
data Module = Module
    { module_name               :: ModuleName
    , module_imports            :: [ImportDecl]
    , module_binds              :: [Binding]
    , module_mb_orch            :: Maybe OrchDict
    }
  deriving (Eq,Show)

data ImportDecl = ImportTop     ModuleName
                | ImportAlias   ModuleName Ident
  deriving (Eq,Show)

type OrchDict = IntMap.IntMap BindingRhs


data FormExpr = PrimForm [Binding]
              | FormName Ident
              | FormExtension FormExpr FormExpr
  deriving (Eq,Show)


data Binding = Binding Ident BindingRhs
  deriving (Eq,Show)

data BindingRhs = BindsM [Ident] DoBlock 
                | BindsP [Ident] Expr
                | BindsF FormExpr
  deriving (Eq,Show)


type DoBlock = [Stmt]

-- | In TopSyntax lambdas are /statements/ rather than 
-- expressions.
--
-- This is to keep the separation between do-stmts and pure 
-- expressions.
--
data Stmt = Return    Expr
          | PrimCall  Ident     [Expr]
          | DoBind    VarBind   DoBlock
          | DoLet     Decl
          | DoLam     Ident     [Ident]   DoBlock
          | DoIf      Expr      DoBlock   DoBlock
  deriving (Eq,Show)


data VarBind = Bind1     Ident
             | BindTuple [Ident]
  deriving (Eq,Show)

-- | Decls only support binding simple values (cannot bind 
-- functions).
--
data Decl = Decl Ident Expr
  deriving (Eq,Show)


-- | Expressions are pure (?)
--
-- Applications are builtin Csound (pure) functions like abs.
--
-- There are no pfields in the top level syntax - pfields are
-- positional and hence not modular. We synthesize them from the 
-- orchestra spec later in compilation.
-- 
data Expr = Lit     Literal
          | Var     Ident
          | FormRef Ident
          | App     Ident     [Expr]
          | Cond    Expr      Expr      Expr
          | UnaryE  UnaryOp   Expr
          | BinE    BinOp     Expr      Expr
          | RelE    RelOp     Expr      Expr
          | Tuple   [Expr]
  deriving (Eq,Show)


data Literal = Bool     Bool
             | Int      Int
             | Float    Decimal
             | String   String
  deriving (Eq,Show)


--------------------------------------------------------------------------------
-- Pretty printing


ppProgram :: Program -> Doc
ppProgram (Program { prog_main         = main_mod
                   , prog_extra_libs   = modules }) = 
    vspaced $ ppModule main_mod : map ppModule modules


ppModule :: Module -> Doc
ppModule (Module { module_name      = name
                 , module_imports   = imports
                 , module_binds     = binds
                 , module_mb_orch   = mb_orch }) = 
    let modu  = text "module" <+> ppModuleName name <+> text "where"
        imps  = vspaced $ map ppImportDecl imports 
        orch  = case mb_orch of { Nothing -> empty; Just o -> ppOrchDict o }
        bdocs = orch : map ppBinding binds
    in vspaced $ modu : imps : bdocs


ppImportDecl :: ImportDecl -> Doc
ppImportDecl (ImportTop s)      = text "import" <+> ppModuleName s
ppImportDecl (ImportAlias s t)  = 
    text "import" <+> ppModuleName s <+> text "as" <+> ppIdent t

ppOrchDict :: OrchDict -> Doc
ppOrchDict dict = text "orch" <+> char '=' 
    $+$ nest 2 (vcommaBraces $ map fn $ IntMap.toAscList dict)
  where
    fn (ix,rhs) = int ix <+> char '=' <+> ppBindingRhs rhs

ppBinding :: Binding -> Doc
ppBinding (Binding vid body) 
    | isForm body = text "def" <+> ppIdent vid <+> char '=' 
                      $+$ nest 2 (ppBindingRhs body)
    | otherwise   = ppIdent vid <+> char '=' <+> ppBindingRhs body

isForm :: BindingRhs -> Bool
isForm (BindsF {}) = True
isForm _           = False


ppFormExpr :: FormExpr -> Doc
ppFormExpr (PrimForm [])        = text "{}"
ppFormExpr (PrimForm xs)        = vcommaBraces (map ppBinding xs)
ppFormExpr (FormName vid)       = ppIdent vid
ppFormExpr (FormExtension f1 f2) = ppFormExpr f1 <+> char '@' <+> ppFormExpr f2


ppBindingRhs :: BindingRhs -> Doc
ppBindingRhs (BindsM args body) = methodDefnBody args (ppDoBlock body)
ppBindingRhs (BindsP args body) = methodDefnBody args (ppExpr body)
ppBindingRhs (BindsF body)      = ppFormExpr body


methodDefnBody :: [Ident] -> Doc -> Doc
methodDefnBody []   body = body
methodDefnBody args body = let formals = hsep $ map ppIdent args
                           in char '\\' <> formals <+> text "->" <+> body


ppDecl :: Decl -> Doc
ppDecl (Decl vid e)             = ppIdent vid <+> char '=' <+> ppExpr e


ppDoBlock :: DoBlock -> Doc
ppDoBlock [x]                   = ppStmt x
ppDoBlock xs                    = text "do" <+> vsemiBraces (map ppStmt xs)

ppStmt :: Stmt -> Doc
ppStmt (Return e)               = text "return" <+> ppExpr e
ppStmt (PrimCall v es)          = ppIdent v <+> commaSep (map ppExpr es)
ppStmt (DoBind a de)            = ppVarBind a <+> text "<-" <+> ppDoBlock de
ppStmt (DoLet d)                = text "let" <+> ppDecl d
ppStmt (DoLam s args de)        = 
    text "letfun" <+> ppIdent s <+> hsep (map ppIdent args) 
                  <+> char '=' <+> ppDoBlock de 
ppStmt (DoIf c d1 d2)           = 
    text "if" <+> ppExpr c <+> nest 2 (    text "then" <+> ppDoBlock d1
                                       $+$ text "else" <+> ppDoBlock d2)


ppVarBind :: VarBind -> Doc
ppVarBind (Bind1 v)             = ppIdent v
ppVarBind (BindTuple vs)        = parens (commaSep $ map ppIdent vs)


ppExpr :: Expr -> Doc
ppExpr = unparse . exprDoc

exprDoc :: Expr -> DocE
exprDoc (Lit v)                 = Atom $ ppLiteral v
exprDoc (Var s)                 = Atom $ ppIdent s
exprDoc (FormRef s)             = Atom $ ppIdent s 
exprDoc (App vid es)            = 
    Atom $ ppIdent vid <> parens (commaSep $ map ppExpr es)

exprDoc (Cond c e1 e2)          = 
    Atom $ text "cond" <+> ppExpr c <+> char ':' <+> ppExpr e1
                                    <+> char '?' <+> ppExpr e2

exprDoc (UnaryE op e)           = unaryDocE op (exprDoc e)
exprDoc (BinE op e1 e2)         = binDocE op (exprDoc e1) (exprDoc e2)
exprDoc (RelE relop e1 e2)      = relDocE relop (exprDoc e1) (exprDoc e2)
exprDoc (Tuple es)              = Atom $ tupled (map ppExpr es) 


ppLiteral :: Literal -> Doc
ppLiteral (Bool b)              = text $ show b
ppLiteral (Int i)               = int i
ppLiteral (Float d)             = text $ truncatedDouble $ realToFrac d
ppLiteral (String s)            = doubleQuotes $ text s




--------------------------------------------------------------------------------
-- Operations

moduleDict :: [Module] -> ModuleDict
moduleDict = foldr fn Map.empty
  where
    fn a = Map.insert (module_name a) a

findModule :: ModuleName -> ModuleDict -> Maybe Module
findModule = Map.lookup 